import typing

import pandas as pd

from pyspark.sql.types import StructType, StructField, LongType, StringType
from pyspark.sql.streaming.state import GroupStateTimeout, GroupStateImpl

output_type = StructType([
    StructField("key", LongType()),
    StructField("countAsString", StringType())])
state_type = StructType([StructField("count", LongType())])

# Type hints are optional in `func`.
def func(key: typing.Tuple, pdf: pd.DataFrame, state: GroupStateImpl) -> pd.DataFrame:
    count = state.getOption
    if count is None:
        count = 0
    else:
        count = count[0]
    count += len(pdf)
    state.update((count,))
    return pd.DataFrame({'key': [key[0]], 'countAsString': [str(count)]})


df = spark.readStream.format("rate").option("rowsPerSecond", 10).load().selectExpr("value % 3 as v")
df.groupBy(df["v"]).applyInPandasWithState(
    func, output_type, state_type, "Update", GroupStateTimeout.NoTimeout
).writeStream.format("console").queryName("test").start()